//============================================================================
// Name        : slicer.c
// Author      : Yugi
// Version     : 1.0
// Description : 
//============================================================================

#include "slicer.h"
#include "parameter.h"

/*
* Generate an initial slice that contains only the target function instructions.
*
* RETURN:	true if initial slice is complete, false otherwise.
* MODIFY:	iList
*			pcfg
*			target
*/
bool init_slice(void *drcontext, instrlist_t *bb, app_pc *target, PCFG pcfg, 
				instrlist_t *iList){
	static bool fSlice = false;		// Flag for slicing

	/* First instruction of this basic block */
	instr_t *pFirst = instrlist_first(bb);
	/* Last instruction of this basic block */
	instr_t *pLast = instrlist_last(bb);

	/* Target hit */
	if (*target == instr_get_app_pc(pFirst))
		fSlice = true;

	/* Add this basic block to cfg */
	if (fSlice){
		/* iterate through pcfg->basic_blocks */
		instr_t *pInstr = pFirst;

		/* For every instruction, clone & append to iList */
		while (pInstr){
			instrlist_append(iList, instr_clone(drcontext, pInstr));
			pInstr = instr_get_next(pInstr);
		}

		cfg_add_bb(drcontext, pcfg, instr_get_app_pc(pFirst));

		if (instr_is_call(pLast)){
			/* Return address of the call instruction */
			*target = instr_get_app_pc(pLast) + instr_length(drcontext, pLast);
			fSlice = false;

			return false;
		}else if (instr_is_return(pLast)){
			/* Reach a return instruction, the initial slice is complete */
			fSlice = false;
			return true;
		}
	}

	return false;
}

/*
* RETURN:	instructions count of iList
*/
unsigned int count(instrlist_t *iList){
	instr_t *pInstr;		// instruction iterator
	unsigned int iCount;	// Count of insrtuctions in iList

	pInstr = instrlist_first(iList);
	iCount = 0;

	while (pInstr){
		iCount++;
		pInstr = instr_get_next(pInstr);
	}

	return iCount;
}

/*
* Dynamic backward slicing
*
* REQUIRE:	iList not NULL
*			worklist not NULL
* RETURN:	bitmap that marks instruction
* MODIFY:	worklist
*/
PSBITMAP dynamic_backward_slice(instrlist_t *iList, PWORKLIST worklist){
	PSBITMAP bitmap;		// bitmaps for marking instructions
	unsigned int index;		// instruction index
	instr_t *pInstr;		// instruction iterator
	bool flag = false;		// Whether we mark the instruction or not
	unsigned int iCount;	// Count of insrtuctions in iList
	unsigned int opCount, opIndex;	// Operand count and index

	/* init bitmaps */
	iCount = count(iList);
	bitmap = bm_new(iCount);

	/* backward iterate */
	pInstr = instrlist_last(iList);
	index = iCount - 1;

	while (pInstr){
		flag = false;
		opCount = instr_num_dsts(pInstr);
		opIndex = 0;

		/* Iterate through dst-operands in 'pInstr' */
		while (opIndex < opCount){
			/* Remove this operand from worklist */
			flag = flag || wl_kill(worklist, instr_get_dst(pInstr, opIndex));
			opIndex++;
		}

		if (flag){
			opCount = instr_num_srcs(pInstr);
			opIndex = 0;

			/* Iterate through src-operands in 'pInstr' */
			while (opIndex < opCount){
				/* Insert into worklist */
				wl_gen(worklist, instr_get_src(pInstr, opIndex));
				opIndex++;
			}

			bm_setbit(bitmap, index);
		}

		pInstr = instr_get_prev(pInstr);
		index--;
	}

	return bitmap;
}

/*
* Dynamic forward slicing
*
* REQUIRE:	worklist not NULL
*			iList not NULL
*			iCount is greater than 0
* RETURN:	bitmap that marks instruction
* MODIFY:	worklist
*/
PSBITMAP dynamic_forward_slice(instrlist_t *iList, PWORKLIST worklist){
	PSBITMAP bitmap;		// bitmaps for marking instructions
	unsigned int index;		// instruction index
	instr_t *pInstr;		// instruction iterator
	bool flag = false;		// Whether we mark the instruction or not
	unsigned int iCount;	// Count of insrtuctions in iList
	unsigned int opCount, opIndex;	// Operand count and index

	/* init bitmaps */
	iCount = count(iList);
	bitmap = bm_new(iCount);

	/* backward iterate */
	pInstr = instrlist_first(iList);
	index = 0;

	while (pInstr){
		flag = false;

		if (instr_num_dsts(pInstr)){
			opCount = instr_num_srcs(pInstr);
			opIndex = 0;

			/* Iterate through src-operands in 'pInstr' */
			while (opIndex < opCount){
				/* Test if worklist contains this operand */
				if (wl_contain(worklist, instr_get_src(pInstr, opIndex))){
					flag = true;
					break;
				}

				opIndex++;
			}

			opCount = instr_num_dsts(pInstr);
			opIndex = 0;

			/* Iterate through dst-operands in 'pInstr' */
			while (opIndex < opCount){
				wl_kill(worklist, instr_get_dst(pInstr, opIndex));
				opIndex++;
			}

			if (flag){
				opIndex = 0;

				/* Iterate through dst-operands in 'pInstr' */
				while (opIndex < opCount){
					/* Insert into worklist */
					wl_gen(worklist, instr_get_dst(pInstr, opIndex));
					opIndex++;
				}

				bm_setbit(bitmap, index);
			}
		}

		pInstr = instr_get_next(pInstr);
		index++;
	}

	return bitmap;
}

/*
* Static backward slicing
*
* REQUIRE:	worklist not NULL
*			pcfg not NULL
* MODIFY:	pcfg
*/
void static_backward_slice(void *drcontext, PCFG pcfg, PWORKLIST worklist){
	unsigned int index;					// index for pcfg->nodes
	unsigned int exit;					// exit point of pcfg->nodes
	PEDGE edge;							// Edge iterator
	PINT_STACK is, is_temp, is_itr;		// Stack and stack temporary pointer
	PSBITMAP bm, bm_temp;				// Bitmap and bitmap temporary pointer
	PWORKLIST wl, wl_temp;				// Worklist and worklist temporary pointer

	/* iterate through pcfg->nodes */
	for (index = 0; index<pcfg->size; index++){
		pcfg->nodes[index].pList = NULL;
		pcfg->nodes[index].bitmap = bm_new(count(pcfg->nodes[index].basic_block));
		pcfg->nodes[index].visited = false;
	}
	
	/* Find out exit point of pcfg->nodes */
	for (exit = 0; exit<pcfg->size; exit++){
		bool flag = true;
		edge = pcfg->head;

		while (edge){
			if (edge->src == exit){
				flag = false;
				break;
			}

			edge = edge->next;
		}

		if (flag)
			break;
	}

	/* Incorrect cfg */
	if (exit == pcfg->size)
		return;

	/* Static slicing */
	pcfg->nodes[exit].pList = wl_clone(worklist);

	/* Push to stack */
	is = (PINT_STACK) malloc(sizeof(INT_STACK));
	is->index = exit;
	is->previous = NULL;
	pcfg->stack = is;

	/* Loop until stack being NULL */
	while (pcfg->stack){
		/* Pop stack */
		is = pcfg->stack;
		pcfg->stack = is->previous;
		wl = wl_clone(pcfg->nodes[is->index].pList);

		/* Backward slice */
		bm = dynamic_backward_slice(pcfg->nodes[is->index].basic_block, wl);
		/* Mark visited */
		pcfg->nodes[is->index].visited = true;

		/* Merge bitmap */
		if (pcfg->nodes[is->index].bitmap){
			bm_temp = bm_merge(pcfg->nodes[is->index].bitmap, bm);
			bm_delete(bm);
			bm_delete(pcfg->nodes[is->index].bitmap);
			pcfg->nodes[is->index].bitmap = bm_temp;
		}else{
			pcfg->nodes[is->index].bitmap = bm;
		}

		/* Get current's previous nodes and push them to stack */
		edge = pcfg->head;

		while (edge){
			if (edge->dst == is->index){
				is_temp = (PINT_STACK) malloc(sizeof(INT_STACK));
				is_temp->index = edge->src;

				is_itr = pcfg->stack;

				/* Stack already has this node, remove it */
				while (is_itr){
					if (is_itr->index == is_temp->index){
						if (is_itr->previous){
							is_itr->index = is_itr->previous->index;
							is_itr->previous = is_itr->previous->previous;
						}else{
							is_itr = NULL;
						}

						break;
					}

					is_itr = is_itr->previous;
				}

				/* 
				* This algorithom uses worklist to determine whether the next node needs to be
				* pushed into stack. If the target node is visited and the current worklist is 
				* subset of the previous worklist, then we won't push it into stack.
				*/
				if (pcfg->nodes[is_temp->index].pList){
					if (pcfg->nodes[is_temp->index].visited && !wl_subset(pcfg->nodes[is_temp->index].pList, wl)){
						edge = edge->next;
						continue;
					}
				}

				/* Merge worklist. */
				if (wl->size > 0)
					if (pcfg->nodes[is_temp->index].pList){
						wl_temp = wl_merge(pcfg->nodes[is_temp->index].pList, wl);
						/* 
						* NOTES: Mightbe a bug here.
						* If we delete the worklist, it will certainly free the space held by
						* worklist in heap. However, the system will cause a segment fault when
						* malloc after quite a lot malloc & free success.
						* If you want to reconstruct this bug, just remove the // in the next
						* statement. It will block in wl_new() and to be specific, it is block
						* in "PWORKLIST pList = (PWORKLIST) malloc(sizeof(WORKLIST));"
						* Use testcase1.exe if other input makes no sense.
						*
						* Debug it further.
						*/
						//wl_delete(pcfg->nodes[is_temp->index].pList);
						pcfg->nodes[is_temp->index].pList = wl_temp;
					}else{
						pcfg->nodes[is_temp->index].pList = wl_clone(wl);
					}

					/* Push to stack */
				is_temp->previous = pcfg->stack;
				pcfg->stack = is_temp;
			}

			edge = edge->next;
		}

		wl_delete(wl);
	}
}

/*
* Static forward slicing
*
* REQUIRE:	worklist not NULL
*			pcfg not NULL
* MODIFY:	pcfg
*/
void static_forward_slice(void *drcontext, PCFG pcfg, PWORKLIST worklist){
	unsigned int index;					// index for pcfg->nodes
	unsigned int entry;					// entry point of pcfg->nodes
	PEDGE edge;							// Edge iterator
	PINT_STACK is, is_temp, is_itr;		// Stack and stack temporary pointer
	PSBITMAP bm, bm_temp;				// Bitmap and bitmap temporary pointer
	PWORKLIST wl, wl_temp;				// Worklist and worklist temporary pointer

	/* iterate through pcfg->nodes */
	for (index = 0; index<pcfg->size; index++){
		pcfg->nodes[index].pList = NULL;
		pcfg->nodes[index].bitmap = bm_new(count(pcfg->nodes[index].basic_block));
		pcfg->nodes[index].visited = false;
	}
	
	/* Find out exit point of pcfg->nodes */
	for (entry = 0; entry<pcfg->size; entry++){
		bool flag = true;
		edge = pcfg->head;

		while (edge){
			if (edge->dst == entry){
				flag = false;
				break;
			}

			edge = edge->next;
		}

		if (flag)
			break;
	}

	/* Incorrect cfg */
	if (entry == pcfg->size)
		return;

	/* Static slicing */
	pcfg->nodes[entry].pList = wl_clone(worklist);

	/* Push to stack */
	is = (PINT_STACK) malloc(sizeof(INT_STACK));
	is->index = entry;
	is->previous = NULL;
	pcfg->stack = is;

	/* Loop until stack being NULL */
	while (pcfg->stack){
		/* Pop stack */
		is = pcfg->stack;
		pcfg->stack = is->previous;
		wl = wl_clone(pcfg->nodes[is->index].pList);

		/* Forward slice */
		bm = dynamic_forward_slice(pcfg->nodes[is->index].basic_block, wl);
		/* Mark visited */
		pcfg->nodes[is->index].visited = true;

		/* Merge bitmap */
		if (pcfg->nodes[is->index].bitmap){
			bm_temp = bm_merge(pcfg->nodes[is->index].bitmap, bm);
			bm_delete(bm);
			bm_delete(pcfg->nodes[is->index].bitmap);
			pcfg->nodes[is->index].bitmap = bm_temp;
		}else{
			pcfg->nodes[is->index].bitmap = bm;
		}

		/* Get current's previous nodes and push them to stack */
		edge = pcfg->head;

		while (edge){
			if (edge->src == is->index){
				is_temp = (PINT_STACK) malloc(sizeof(INT_STACK));
				is_temp->index = edge->dst;

				is_itr = pcfg->stack;

				/* Stack already has this node, remove it */
				while (is_itr){
					if (is_itr->index == is_temp->index){
						if (is_itr->previous){
							is_itr->index = is_itr->previous->index;
							is_itr->previous = is_itr->previous->previous;
						}else{
							is_itr = NULL;
						}

						break;
					}

					is_itr = is_itr->previous;
				}

				/* 
				* This algorithom uses worklist to determine whether the next node needs to be
				* pushed into stack. If the target node is visited and the current worklist is 
				* subset of the previous worklist, then we won't push it into stack.
				*/
				if (pcfg->nodes[is_temp->index].pList){
					if (pcfg->nodes[is_temp->index].visited && !wl_subset(pcfg->nodes[is_temp->index].pList, wl)){
						edge = edge->next;
						continue;
					}
				}

				/* Merge worklist. */
				if (wl->size > 0)
					if (pcfg->nodes[is_temp->index].pList){
						wl_temp = wl_merge(pcfg->nodes[is_temp->index].pList, wl);
						//wl_delete(pcfg->nodes[is_temp->index].pList);
						pcfg->nodes[is_temp->index].pList = wl_temp;
					}else{
						pcfg->nodes[is_temp->index].pList = wl_clone(wl);
					}

					/* Push to stack */
				is_temp->previous = pcfg->stack;
				pcfg->stack = is_temp;
			}

			edge = edge->next;
		}

		wl_delete(wl);
	}
}

/* 
* Print initial slice
*
* REQUIRE:	iList not NULL
*			outfile not NULL
*/
void print_initial_slice(void *drcontext, file_t outfile, instrlist_t *iList){
	instr_t *pInstr = instrlist_first(iList);

	/* print all the instructions in iList */
	while (pInstr){
		dr_print_instr(drcontext, outfile, pInstr, "");
		pInstr = instr_get_next(pInstr);
	}
}

/*
* Print a slice according to bitmap
*
* REQUIRE:	bitmap not NULL
*			iList not NULL
*			outfile not NULL
*/
void print_dynamic_slice(void *drcontext, file_t outfile, instrlist_t *iList, PSBITMAP bitmap){
	instr_t *pInstr =  instrlist_first(iList);
	unsigned int index = 0;
	
	while (pInstr){
		if (bm_testbit(bitmap, index))
			dr_print_instr(drcontext, outfile, pInstr, "");

		pInstr = instr_get_next(pInstr);
		index++;
	}
}
